#!/usr/bin/env python3
"""
Script to crop and resize preprocessed Hypersim output.

The script reads the preprocessed Hypersim data (each scene folder contains files:
    "xxxxxx_rgb.png", "xxxxxx_depth.npy", "xxxxxx_cam.npz")
and applies a crop and resize operation such that the final output has the given resolution.
The cropping is centered around the principal point (from the camera intrinsics).

Usage:
    python crop_resize_hypersim.py --source_dir /path/to/hypersim_preprocessed \
                                   --target_dir /path/to/cropped_resized_output \
                                   --resolution 800,600
"""

import os
import os.path as osp
import argparse
import shutil
import numpy as np
import PIL.Image
from PIL import Image
import sys

from tqdm import tqdm

current_dir = os.path.dirname(os.path.abspath(__file__))
src_dir = os.path.abspath(os.path.join(current_dir, "..", "src"))
sys.path.insert(0, src_dir)

# Import the cropping utilities (do not modify these implementations)
from utils.cropping import rescale_image_depthmap
import src.dust3r.datasets.utils.cropping as cropping

def crop_resize_if_necessary(image, depthmap, intrinsics, resolution, info=None):
    """
    Crop and resize the given image, depthmap, and intrinsics.
    
    The procedure is as follows:
      1. Ensure the image is a PIL.Image.Image.
      2. Determine a centered crop window on the principal point.
      3. Crop the image and depthmap.
      4. Downscale the result using Lanczos interpolation so that the image size matches resolution.
      5. Compute new intrinsics from the cropping.
      6. Perform final cropping (if necessary) with bilinear interpolation.
    
    Args:
        image: an RGB image (PIL Image or numpy array)
        depthmap: a depth map (numpy array)
        intrinsics: a 3x3 numpy array containing camera intrinsics.
        resolution: a tuple (width, height) specifying target resolution.
        info: optional info for error messages.
    
    Returns:
        A tuple of (image_out, depthmap_out, intrinsics_out)
    """
    # Convert image to PIL image if needed.
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    
    W, H = image.size
    # Use the principal point from intrinsics (assumed to be in position [0,2] and [1,2])
    cx, cy = intrinsics[:2, 2].round().astype(int)
    min_margin_x = min(cx, W - cx)
    min_margin_y = min(cy, H - cy)
    if min_margin_x <= W / 5 or min_margin_y <= H / 5:
        raise ValueError(f"Bad principal point in view {info}")
    # Define crop bounding box centered on (cx, cy)
    l, t = cx - min_margin_x, cy - min_margin_y
    r, b = cx + min_margin_x, cy + min_margin_y
    crop_bbox = (l, t, r, b)
    image, depthmap, intrinsics = cropping.crop_image_depthmap(image, depthmap, intrinsics, crop_bbox)
    
    # Downscale using high-quality Lanczos interpolation so that image.size == resolution.
    target_w = resolution
    target_h = resolution
    image, depthmap, intrinsics = rescale_image_depthmap(image, depthmap, intrinsics, (target_w, target_h))
    W, H = image.size
    resolution = image.size
    
    # Compute new intrinsics from final cropping.
    intrinsics2 = cropping.camera_matrix_of_crop(intrinsics, image.size, resolution, offset_factor=0.5)
    crop_bbox = cropping.bbox_from_intrinsics_in_out(intrinsics, intrinsics2, resolution)
    image, depthmap, intrinsics2 = cropping.crop_image_depthmap(image, depthmap, intrinsics, crop_bbox)
    
    return image, depthmap, intrinsics2

def process_camera_folder(camera_source_dir, camera_target_dir, resolution):
    """
    Process a single camera folder.
    
    The folder is expected to contain files following the pattern:
         "xxxxxx_rgb.png", "xxxxxx_depth.npy", and "xxxxxx_cam.npz"
    
    This function applies the crop and resize operation to each image triplet and saves the results in
    the target folder with the same filenames.
    """
    os.makedirs(camera_target_dir, exist_ok=True)
    # Find all RGB image files in the folder.
    rgb_files = sorted([f for f in os.listdir(camera_source_dir) if f.endswith("_rgb.png")])
    if not rgb_files:
        print(f"No RGB images found in {camera_source_dir}")
        return False
    
    for rgb_file in tqdm(rgb_files, desc=f"Processing {camera_source_dir.split('/')[-1]}"):
        base, ext = osp.splitext(rgb_file)  # e.g., "000000_rgb"
        base_id = base.split('_')[0]  # e.g., "000000"
        rgb_path = osp.join(camera_source_dir, f"{base_id}_rgb.png")
        depth_path = osp.join(camera_source_dir, f"{base_id}_depth.npy")
        cam_path = osp.join(camera_source_dir, f"{base_id}_cam.npz")
        
        if not (osp.isfile(rgb_path) and osp.isfile(depth_path) and osp.isfile(cam_path)):
            print(f"Missing files for base id {base_id} in {camera_source_dir}, skipping.")
            return False
        
        try:
            image = Image.open(rgb_path)
        except Exception as e:
            print(f"Error loading image {rgb_path}: {e}")
            return False
        
        try:
            depthmap = np.load(depth_path)
        except Exception as e:
            print(f"Error loading depth map {depth_path}: {e}")
            return False
        
        try:
            cam_data = np.load(cam_path)
            intrinsics = cam_data["intrinsics"]
            pose = cam_data["pose"]
        except Exception as e:
            print(f"Error loading camera metadata {cam_path}: {e}")
            return False
        
        try:
            image_cropped, depth_cropped, new_intrinsics = crop_resize_if_necessary(
                image, depthmap, intrinsics, resolution, info=base_id
            )
        except Exception as e:
            print(f"Error processing {base_id}: {e}")
            return False
        
        # Save the processed files in the target folder.
        out_rgb = osp.join(camera_target_dir, f"{base_id}_rgb.png")
        out_depth = osp.join(camera_target_dir, f"{base_id}_depth.npy")
        out_cam = osp.join(camera_target_dir, f"{base_id}_cam.npz")
        
        try:
            image_cropped.save(out_rgb)
        except Exception as e:
            print(f"Error saving image {out_rgb}: {e}")
            return False
        try:
            np.save(out_depth, depth_cropped)
        except Exception as e:
            print(f"Error saving depth map {out_depth}: {e}")
            return False
        try:
            np.savez(out_cam, intrinsics=new_intrinsics, pose=pose)
        except Exception as e:
            print(f"Error saving camera metadata {out_cam}: {e}")
            return False
    return True


def process_scene(scene_source_dir, scene_target_dir, resolution):
    """
    Process a single scene folder.
    
    Each scene folder (e.g., "ai_022_005") is expected to contain one or more camera subfolders
    (e.g., "cam_00", "cam_01", etc.). This function processes each camera folder with the
    crop/resize operation and writes the outputs into the corresponding subfolder in the target directory.
    """
    # Enumerate camera subfolders.
    camera_folders = sorted([d for d in os.listdir(scene_source_dir) if osp.isdir(osp.join(scene_source_dir, d))])
    if not camera_folders:
        print(f"No camera subfolders found in {scene_source_dir}")
        return False
    
    os.makedirs(scene_target_dir, exist_ok=True)
    for cam_folder in tqdm(camera_folders, desc=f"Processing {scene_source_dir.split('/')[-1]}"):
        camera_source_dir = osp.join(scene_source_dir, cam_folder)
        camera_target_dir = osp.join(scene_target_dir, cam_folder)
        # print(f"Processing camera folder {cam_folder} ...")
        success = process_camera_folder(camera_source_dir, camera_target_dir, resolution)
        if not success:
            return False
    return True


def main():
    parser = argparse.ArgumentParser(
        description="Crop and resize preprocessed Hypersim output while preserving the scene/camera folder structure."
    )
    parser.add_argument("--source_dir", default="/lc/data/3D/processed_hypersim/test",
                        help="Source directory of preprocessed Hypersim (e.g., processed_hypersim)")
    parser.add_argument("--target_dir", default="/lc/data/3D/resized_hypersim/test",
                        help="Target directory to save cropped/resized output")
    parser.add_argument("--resolution", default=512, type=int, help="images resolution")
    
    args = parser.parse_args()
    
    # The source directory (e.g., processed_hypersim) contains scene subfolders.
    scenes = sorted([d for d in os.listdir(args.source_dir) if osp.isdir(osp.join(args.source_dir, d))])
    if not scenes:
        print("No scene subdirectories found in the source directory.")
        sys.exit(1)
    
    for scene in tqdm(scenes, desc=f"Processing scenes"):
        scene_source_dir = osp.join(args.source_dir, scene)
        scene_target_dir = osp.join(args.target_dir, scene)
        # print(f"Processing scene {scene} ...")
        success = process_scene(scene_source_dir, scene_target_dir, args.resolution)
        if not success:
            print(f"Error processing scene {scene}")
            # return
        else:
            shutil.rmtree(scene_source_dir)
        # break
    print("Done processing all scenes.")


if __name__ == "__main__":
    main()